{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 变分量子态对角化算法\n",
    "\n",
    "<em> Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved. </em>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 概览\n",
    "\n",
    "- 在本案例中，我们将通过 Paddle Quantum 训练量子神经网络来完成量子态的对角化\n",
    "\n",
    "- 首先，让我们通过下面几行代码引入必要的 library 和 package。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:17:03.603558Z",
     "start_time": "2021-04-30T09:17:00.931471Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy\n",
    "from numpy import diag\n",
    "from numpy import pi as PI\n",
    "import scipy\n",
    "import scipy.stats\n",
    "import paddle\n",
    "from paddle import matmul, trace\n",
    "import paddle_quantum\n",
    "from paddle_quantum import intrinsic\n",
    "from paddle_quantum.ansatz import Circuit\n",
    "from paddle_quantum.linalg import dagger"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 背景\n",
    "量子态对角化算法（variational quantum state diagonalization, VQSD）[1-3] 目标是输出一个量子态的本征谱，即其所有本征值。求解量子态的本征值在量子计算中有着诸多应用，比如可以用于计算保真度和冯诺依曼熵，也可以用于主成分分析。\n",
    "- 量子态通常是一个混合态，表示如下 $\\rho_{\\text{mixed}} = \\sum_i P_i |\\psi_i\\rangle\\langle\\psi_i|$\n",
    "- 作为一个简单的例子，我们考虑一个 2 量子位的量子态，它的本征谱为 $(0.5, 0.3, 0.1, 0.1)$，我们先通过随机作用一个酉矩阵来生成具有这样本征谱的随机量子态。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:17:03.622120Z",
     "start_time": "2021-04-30T09:17:03.609542Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.2569-0.j     -0.012 +0.0435j -0.0492-0.0055j -0.0548+0.0682j]\n",
      " [-0.012 -0.0435j  0.2959+0.j      0.1061-0.0713j -0.0392-0.0971j]\n",
      " [-0.0492+0.0055j  0.1061+0.0713j  0.2145-0.j      0.0294-0.1132j]\n",
      " [-0.0548-0.0682j -0.0392+0.0971j  0.0294+0.1132j  0.2327+0.j    ]]\n"
     ]
    }
   ],
   "source": [
    "scipy.random.seed(13)                  # 固定随机种子，方便复现结果\n",
    "V = scipy.stats.unitary_group.rvs(4)   # 随机生成一个酉矩阵\n",
    "D = diag([0.5, 0.3, 0.1, 0.1])         # 输入目标态 rho 的谱\n",
    "V_H = V.conj().T \n",
    "rho = V @ D @ V_H                      # 通过逆向的谱分解生成 rho\n",
    "print(numpy.around(rho, 4))            # 打印量子态 rho"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 搭建量子神经网络\n",
    "- 在这个案例中，我们将通过训练量子神经网络 QNN（也可以理解为参数化量子电路）来学习量子态的本征谱。这里，我们提供一个预设的 2 量子位量子电路。\n",
    "\n",
    "- 我们预设一些该参数化电路的参数，比如宽度为 2 量子位。\n",
    "\n",
    "- 初始化其中的变量参数，${\\bf{\\theta }}$ 代表我们量子神经网络中的参数组成的向量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:17:03.646986Z",
     "start_time": "2021-04-30T09:17:03.637134Z"
    }
   },
   "outputs": [],
   "source": [
    "N = 2            # 量子神经网络的宽度\n",
    "SEED = 14        # 固定随机种子\n",
    "THETA_SIZE = 15  # 量子神经网络中参数的数量\n",
    "\n",
    "def U_theta(N: int) -> Circuit:\n",
    "    \"\"\"\n",
    "    Quantum Neural Network\n",
    "    \"\"\"\n",
    "    \n",
    "    # 按照量子比特数量/网络宽度初始化量子神经网络\n",
    "    cir = Circuit(N)\n",
    "    # 调用内置的量子神经网络模板\n",
    "    cir.universal_two_qubits([0, 1])\n",
    "    # 返回量子神经网络的电路\n",
    "    return cir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 配置训练模型——损失函数\n",
    "- 现在我们已经有了数据和量子神经网络的架构，我们将进一步定义训练参数、模型和损失函数。\n",
    "- 通过作用量子神经网络 $U(\\theta)$ 在量子态 $\\rho$ 后得到的量子态记为 $\\tilde\\rho$，我们设定损失函数为 $\\tilde\\rho$ 与用来标记的量子态 $\\sigma=0.1 |00\\rangle\\langle 00| + 0.2 |01\\rangle \\langle 01| + 0.3 |10\\rangle \\langle10| + 0.4 |11 \\rangle\\langle 11|$ 的内积。\n",
    "- 具体的，设定损失函数为 $\\mathcal{L}(\\boldsymbol{\\theta})  = \\text{Tr}(\\tilde\\rho\\sigma)$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:17:03.676479Z",
     "start_time": "2021-04-30T09:17:03.660557Z"
    }
   },
   "outputs": [],
   "source": [
    "# 输入用来标记的量子态sigma\n",
    "sigma = diag([0.1, 0.2, 0.3, 0.4])\n",
    "\n",
    "class Net(paddle.nn.Layer):\n",
    "    \"\"\"\n",
    "    Construct the model net\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, rho, sigma):\n",
    "        super(Net, self).__init__()\n",
    "        \n",
    "        # 将 numpy.ndarray 转换成 Paddle 中支持的 Tensor\n",
    "        complex_dtype = paddle_quantum.get_dtype()\n",
    "        self.rho = paddle.cast(paddle.to_tensor(rho), complex_dtype)\n",
    "        self.sigma = paddle.cast(paddle.to_tensor(sigma), complex_dtype)\n",
    "\n",
    "        # 施加量子神经网络\n",
    "        self.cir = U_theta(N)                                  \n",
    "\n",
    "    # 定义损失函数和前向传播机制\n",
    "    def forward(self, N):\n",
    "            \n",
    "        # rho_tilde 是将 U 作用在 rho 后得到的量子态 U*rho*U^dagger \n",
    "        U = self.cir.unitary_matrix()\n",
    "        rho_tilde = matmul(matmul(U, self.rho), dagger(U))\n",
    "\n",
    "        # 计算损失函数\n",
    "        loss = trace(matmul(self.sigma, rho_tilde))\n",
    "\n",
    "        return paddle.real(loss), rho_tilde, self.cir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 配置训练模型——模型参数\n",
    "\n",
    "在进行量子神经网络的训练之前，我们还需要进行一些训练的超参数设置，主要是学习速率 (learning rate, LR)和迭代次数(iteration, ITR)。这里我们设定学习速率为 0.1，迭代次数为 50 次。读者不妨自行调整来直观感受下超参数调整对训练效果的影响。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:17:04.114466Z",
     "start_time": "2021-04-30T09:17:04.099242Z"
    }
   },
   "outputs": [],
   "source": [
    "ITR = 50  # 设置训练的总的迭代次数\n",
    "LR = 0.1  # 设置学习速率"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 进行训练\n",
    "\n",
    "- 当训练模型的各项参数都设置完成后，我们将数据转化为 Paddle 中的张量，进而进行量子神经网络的训练。\n",
    "- 过程中我们用的是 Adam Optimizer，也可以调用 Paddle 中提供的其他优化器。\n",
    "- 我们将训练过程中的结果依次输出。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:17:10.042819Z",
     "start_time": "2021-04-30T09:17:08.580527Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter: 0 loss: 0.2494\n",
      "iter: 10 loss: 0.1959\n",
      "iter: 20 loss: 0.1843\n",
      "iter: 30 loss: 0.1816\n",
      "iter: 40 loss: 0.1805\n"
     ]
    }
   ],
   "source": [
    "paddle.seed(SEED)\n",
    "\n",
    "# 确定网络的参数维度\n",
    "net = Net(rho=rho, sigma=sigma)\n",
    "\n",
    "# 一般来说，我们利用 Adam 优化器来获得相对好的收敛\n",
    "# 当然你可以改成 SGD 或者是 RMS prop.\n",
    "opt = paddle.optimizer.Adam(learning_rate=LR, parameters=net.parameters())\n",
    "\n",
    "# 优化循环\n",
    "for itr in range(ITR):\n",
    "\n",
    "    # 前向传播计算损失函数并返回估计的能谱\n",
    "    loss, rho_tilde, cir = net(N)\n",
    "    rho_tilde_np = rho_tilde.numpy()\n",
    "\n",
    "    # 反向传播极小化损失函数\n",
    "    loss.backward()\n",
    "    opt.minimize(loss)\n",
    "    opt.clear_grad()\n",
    "\n",
    "    # 打印训练结果\n",
    "    if itr % 10 == 0:\n",
    "        print('iter:', itr, 'loss:', '%.4f' % loss.numpy()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 总结\n",
    "\n",
    "根据上面训练得到的结果，通过大概50次迭代，我们就比较好的完成了对角化。我们可以通过打印 $\\tilde{\\rho} = U(\\boldsymbol{\\theta})\\rho U^\\dagger(\\boldsymbol{\\theta})$ 的来验证谱分解的效果。特别的，我们可以验证它的对角线与目标谱非常接近。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:17:13.536415Z",
     "start_time": "2021-04-30T09:17:13.527756Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The estimated spectrum is: [0.49938068 0.29916352 0.10103808 0.10041767]\n",
      "The target spectrum is: [0.5 0.3 0.1 0.1]\n"
     ]
    }
   ],
   "source": [
    "print(\"The estimated spectrum is:\", numpy.real(numpy.diag(rho_tilde_np)))\n",
    "print(\"The target spectrum is:\", numpy.diag(D))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_______\n",
    "\n",
    "## 参考文献\n",
    "\n",
    "[1] Larose, R., Tikku, A., Neel-judy, É. O., Cincio, L. & Coles, P. J. Variational quantum state diagonalization. [npj Quantum Inf. (2019) doi:10.1038/s41534-019-0167-6.](https://www.nature.com/articles/s41534-019-0167-6)\n",
    "\n",
    "[2] Nakanishi, K. M., Mitarai, K. & Fujii, K. Subspace-search variational quantum eigensolver for excited states. [Phys. Rev. Res. 1, 033062 (2019).](https://journals.aps.org/prresearch/pdf/10.1103/PhysRevResearch.1.033062)\n",
    "\n",
    "[3] Cerezo, M., Sharma, K., Arrasmith, A. & Coles, P. J. Variational Quantum State Eigensolver. [arXiv:2004.01372 (2020).](https://arxiv.org/pdf/2004.01372.pdf)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
